import sys
import os
import numpy as np
from sklearn.cluster import KMeans
import cv2
import time
import multiprocessing
import snapmodel 
import logging
from utils import *

def gen_refer_mask(img, level):
    refer_mask = np.zeros(img.shape, dtype=np.uint8)
    mult = 4
    if level == 1:
        refer_mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(255)
    elif level == 2:
        refer_mask[8*mult:40*mult, 2*mult:(C-2*mult+1)] = int(255)
    elif level == 3:
        refer_mask[10*mult:40*mult, 5*mult:(C-5*mult+1)] = int(255)
    else:
        refer_mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(255)
    return refer_mask

def datasets(path):
    imgfiles = [os.path.join(path, f) for f in os.listdir(path) if f[-9:].lower() == 'cntnt.jpg']
    return imgfiles

def predict(img, refer_mask):
    focus_mask, scan_mask, raw_mask = snapmodel.make_mask(img, refer_mask)
    return focus_mask, scan_mask, raw_mask

def evaluate(img, mask, refer_mask):
    mask_inside_cost, mask_outside_cost = 0, 0

    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #sgray = better_image(gray, 'ScaleAbs')
    #n1gray = better_image(gray, 'NormalizeL1')
    #n2gray = better_image(gray, 'NormalizeL2')
    #n3gray = better_image(gray, 'NormalizeMM')
    egray = snapmodel.better_image(gray, 'EqualizeHist', refer_mask)
    #agray = better_image(gray, 'AdaptHist')

    #cv2.imshow('gray', gray)
    #cv2.imshow('sgray', sgray)
    #cv2.imshow('n1gray', n1gray)
    #cv2.imshow('n2gray', n2gray)
    #cv2.imshow('n3gray', n3gray)
    #cv2.imshow('egray', egray)
    #cv2.imshow('agray', agray)
    gray = egray

    # 腐蚀
    #kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    #kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3,3))
    #erode_mask = cv2.erode(mask, kernel)

    # 膨胀
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (10,10))
    #kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (15,15))
    dilate_mask = cv2.dilate(mask, kernel)

    inside = np.minimum(gray, mask)
    inside = inside - np.zeros(mask.shape, dtype=np.uint8)
    #mask_inside_cost = np.linalg.norm(inside - np.zeros(mask.shape, dtype=np.uint8))
    if len(inside.nonzero()[0]) > 0:
        mask_inside_cost = np.linalg.norm(inside) / len(inside.nonzero()[0])

    outside = np.maximum(gray, np.ones(mask.shape, dtype=np.uint8) * 255 - (dilate_mask - mask))
    outside = np.ones(mask.shape, dtype=np.uint8) * 255 - outside
    #mask_outside_cost = np.linalg.norm(np.ones(mask.shape, dtype=np.uint8) * 255 - outside)
    if len(outside.nonzero()[0]) > 0:
        mask_outside_cost = np.linalg.norm(outside) / len(outside.nonzero()[0])

    #cv2.imshow('inside', inside)
    #cv2.imshow('outside', outside)
    #cv2.imshow('mask', mask)
    #cv2.imshow('erode_mask', erode_mask)
    #cv2.imshow('dilate_mask', dilate_mask)
    #cv2.waitKey()
    #cv2.destroyAllWindows()

    return mask_inside_cost, mask_outside_cost

def proc(imgfile):
    img = snapmodel.load_image(imgfile)

    focus_mask, scan_mask, raw_mask = predict(img, gen_refer_mask(img, 1))
    mask = raw_mask
    mask_inside_cost, mask_outside_cost = evaluate(img, mask * 255, gen_refer_mask(mask, 1))

    if mask_inside_cost > 1.0:
        focus_mask, scan_mask, raw_mask = predict(img, gen_refer_mask(img, 3))
        mask = raw_mask
        mask_inside_cost, mask_outside_cost = evaluate(img, mask * 255, gen_refer_mask(mask, 3))
    elif mask_inside_cost > 0.8:
        focus_mask, scan_mask, raw_mask = predict(img, gen_refer_mask(img, 2))
        mask = raw_mask
        mask_inside_cost, mask_outside_cost = evaluate(img, mask * 255, gen_refer_mask(mask, 2))

    filename = os.path.basename(imgfile)
    head, tail = filename.split('.')
    #cv2.imwrite(f'{head}_mask.{tail}', (focus_mask + scan_mask) * 127)
    cv2.imwrite(f'{head}_mask.{tail}', (raw_mask) * 127)
    cv2.imwrite(f'{head}.{tail}', img)

    logging.info('file: {}, insidecost: {}, outsidecost: {}'.format(os.path.basename(imgfile), mask_inside_cost, mask_outside_cost))

if __name__=='__main__':
    imgfiles = datasets('SnapshotDat/cntnt_new')

    try:
        os.remove('1.log')
    except:
        pass
    logging.basicConfig(level=1, format='%(message)s', filename=r'1.log')

    pool = multiprocessing.Pool(8)
    pool.map(proc, imgfiles[:50])
    pool.close()
    pool.join()

